{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune SAM on customized datasets for automatic segmentation\n",
    "\n",
    "SAM is a promptable segmentation methods, but it is simple to adapt SAM for automatic segmentation. \n",
    "\n",
    "The trick is to set the bounding box (xmin, ymin, xmax, ymax) as the image size (0, 0, H, W) for all images.\n",
    "\n",
    "Here is a step-by-step demo for liver segmentation in abdomen CT images.\n",
    "\n",
    "1. Download [dataset](https://drive.google.com/file/d/19XFAth7EjMRDOC6u8Hom7Ep8Z8zUR3-S/view?usp=share_link) and put it to `data/MultiOrganDemo_2D/`\n",
    "2. Start this fine-tuning tutorial\n",
    "\n",
    "Remark: The following two tutorials are bounding box-based segmentation. The bounding box is simulated from the ground truth mask.\n",
    "- finetune_and_inference_tutorial_2D_dataset.ipynb\n",
    "- finetune_and_inference_tutorial_3D_dataset.ipynb\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "join = os.path.join\n",
    "from skimage import io\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import monai\n",
    "from monai.networks import one_hot\n",
    "from segment_anything import SamPredictor, sam_model_registry\n",
    "from segment_anything.utils.transforms import ResizeLongestSide\n",
    "from utils.SurfaceDice import compute_dice_coefficient\n",
    "# set seeds\n",
    "torch.manual_seed(2023)\n",
    "np.random.seed(2023)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_tr_path = 'data/MultiOrganDemo_2D/train/'\n",
    "embed_tr_path = join(data_tr_path, 'img_embed')\n",
    "os.makedirs(embed_tr_path, exist_ok=True)\n",
    "# the dataset contains 13 organs. https://flare22.grand-challenge.org/\n",
    "# Please select the organ you want to segment by setting the label_id.\n",
    "label_id = 1 # liver\n",
    "model_type = 'vit_b'\n",
    "checkpoint = 'work_dir/SAM/sam_vit_b_01ec64.pth'\n",
    "device = 'cuda:0'\n",
    "sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pre-compute image image embedding\n",
    "names = sorted(os.listdir(join(data_tr_path, 'images')))\n",
    "for name in tqdm(names):\n",
    "    image_data = io.imread(join(data_tr_path, 'images', name))\n",
    "    if image_data.shape[-1]>3 and len(image_data.shape)==3:\n",
    "        image_data = image_data[:,:,:3]\n",
    "    if len(image_data.shape)==2:\n",
    "        image_data = np.repeat(image_data[:,:,None], 3, axis=-1)\n",
    "    sam_transform = ResizeLongestSide(sam_model.image_encoder.img_size)\n",
    "    resize_img = sam_transform.apply_image(image_data)\n",
    "    resize_img_tensor = torch.as_tensor(resize_img.transpose(2, 0, 1)).to(device)\n",
    "    input_image = sam_model.preprocess(resize_img_tensor[None,:,:,:]) # (1, 3, 1024, 1024)\n",
    "    assert input_image.shape == (1, 3, sam_model.image_encoder.img_size, sam_model.image_encoder.img_size), 'input image should be resized to 1024*1024'\n",
    "    # pre-compute the image embedding\n",
    "    with torch.no_grad():\n",
    "        embedding = sam_model.image_encoder(input_image)\n",
    "        # save the image embedding\n",
    "        np.save(join(embed_tr_path, name.split('.png')[0]+'.npy'), embedding.cpu().numpy()[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% create a dataset class to load npz data and return back image embeddings and ground truth\n",
    "class DatasetPngNPy(Dataset): \n",
    "    def __init__(self, data_root, label_id=1):\n",
    "        self.data_root = data_root\n",
    "        self.files = sorted(os.listdir(join(self.data_root, 'labels'))) \n",
    "        self.label_id = label_id\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.files)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img_embed = np.load(join(self.data_root, 'img_embed', self.files[index].split('.')[0]+'.npy'))\n",
    "        gt2D = io.imread(join(self.data_root, 'labels', self.files[index]))\n",
    "        gt2D = (gt2D==self.label_id).astype(np.uint8)\n",
    "\n",
    "        # convert img embedding, mask, bounding box to torch tensor\n",
    "        return torch.tensor(img_embed).float(), torch.tensor(gt2D[None, :,:]).long()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% test dataset class and dataloader\n",
    "demo_dataset = DatasetPngNPy(data_tr_path, label_id=label_id)\n",
    "demo_dataloader = DataLoader(demo_dataset, batch_size=8, shuffle=True)\n",
    "for img_embed, gt in demo_dataloader:\n",
    "    # img_embed: (B, 256, 64, 64), gt2D: (B, C, 256, 256)\n",
    "    print(f\"{img_embed.shape=}, {gt.shape=}\")\n",
    "    plt.imshow(gt[5,0,:,:].numpy())\n",
    "    plt.axis('off')\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% set up model for fine-tuning \n",
    "# train data path\n",
    "work_dir = './work_dir'\n",
    "task_name = 'AutoSegDemo2D'\n",
    "# prepare SAM model\n",
    "model_save_path = join(work_dir, task_name)\n",
    "os.makedirs(model_save_path, exist_ok=True)\n",
    "# Set up the optimizer, hyperparameter tuning will improve performance here\n",
    "optimizer = torch.optim.Adam(sam_model.mask_decoder.parameters(), lr=1e-5, weight_decay=0)\n",
    "seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% train\n",
    "num_epochs = 100\n",
    "losses = []\n",
    "best_loss = 1e10\n",
    "train_dataset = DatasetPngNPy(data_tr_path, label_id=label_id)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "sam_model.train()\n",
    "for epoch in range(num_epochs):\n",
    "    epoch_loss = 0\n",
    "    # train\n",
    "    for step, (image_embedding, gt) in enumerate(tqdm(train_dataloader)):\n",
    "        # do not compute gradients for image encoder and prompt encoder\n",
    "        with torch.no_grad():\n",
    "            B,_, H, W = gt.shape\n",
    "            # set the bbox as the image size for fully automatic segmentation\n",
    "            boxes = torch.from_numpy(np.array([[0,0,W,H]]*B)).float().to(device)\n",
    "            sparse_embeddings, dense_embeddings = sam_model.prompt_encoder(\n",
    "                points=None,\n",
    "                boxes=boxes[:, None, :],\n",
    "                masks=None,\n",
    "            )\n",
    "        # predicted masks\n",
    "        mask_predictions, _ = sam_model.mask_decoder(\n",
    "            image_embeddings=image_embedding.to(device), # (B, 256, 64, 64)\n",
    "            image_pe=sam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)\n",
    "            sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)\n",
    "            dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)\n",
    "            multimask_output=False,\n",
    "          )\n",
    "\n",
    "        loss = seg_loss(mask_predictions, gt.to(device))\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        epoch_loss += loss.item()\n",
    "    \n",
    "    epoch_loss /= step\n",
    "    losses.append(epoch_loss)\n",
    "    print(f'EPOCH: {epoch}, Loss: {epoch_loss}')\n",
    "    # save the latest model checkpoint\n",
    "    torch.save(sam_model.state_dict(), join(model_save_path, 'sam_model_latest.pth'))\n",
    "    # save the best model\n",
    "    if epoch_loss < best_loss:\n",
    "        best_loss = epoch_loss\n",
    "        torch.save(sam_model.state_dict(), join(model_save_path, 'sam_model_best.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot loss\n",
    "plt.plot(losses)\n",
    "plt.title('Dice + Cross Entropy Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.show() # comment this line if you are running on a server\n",
    "plt.savefig(join(model_save_path, 'train_loss.png'))\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% compare the segmentation results between the original SAM model and the fine-tuned model\n",
    "# load the original SAM model\n",
    "ori_sam_model = sam_model_registry[model_type](checkpoint=checkpoint).to(device)\n",
    "ori_sam_predictor = SamPredictor(ori_sam_model)\n",
    "\n",
    "ts_img_path = 'data/MultiOrganDemo_2D/test/images'\n",
    "ts_gt_path = 'data/MultiOrganDemo_2D/test/labels'\n",
    "test_names = sorted(os.listdir(ts_img_path))\n",
    "test_names = [name for name in test_names if np.sum(io.imread(join(ts_gt_path, name))==label_id)>0]\n",
    "# random select a test case\n",
    "img_idx = np.random.randint(len(test_names))\n",
    "image_data = io.imread(join(ts_img_path, test_names[img_idx]))\n",
    "if image_data.shape[-1]>3 and len(image_data.shape)==3:\n",
    "    image_data = image_data[:,:,:3]\n",
    "if len(image_data.shape)==2:\n",
    "    image_data = np.repeat(image_data[:,:,None], 3, axis=-1)\n",
    "# read ground truth (gt should have the same name as the image) and simulate a bounding box\n",
    "gt_data = io.imread(join(ts_gt_path, test_names[img_idx]))\n",
    "gt_data = (gt_data==label_id).astype(np.uint8)\n",
    "H, W = gt_data.shape\n",
    "\n",
    "# predict the segmentation mask using the original SAM model\n",
    "ori_sam_predictor.set_image(image_data)\n",
    "ori_sam_seg, _, _ = ori_sam_predictor.predict(point_coords=None, box=np.array([0,0,H,W]), multimask_output=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict the segmentation mask using the fine-tuned model\n",
    "# resize image to 3*1024*1024\n",
    "sam_transform = ResizeLongestSide(sam_model.image_encoder.img_size)\n",
    "resize_img = sam_transform.apply_image(image_data)\n",
    "resize_img_tensor = torch.as_tensor(resize_img.transpose(2, 0, 1)).to(device)\n",
    "input_image = sam_model.preprocess(resize_img_tensor[None,:,:,:]) # (1, 3, 1024, 1024)\n",
    "assert input_image.shape == (1, 3, sam_model.image_encoder.img_size, sam_model.image_encoder.img_size), 'input image should be resized to 1024*1024'\n",
    "\n",
    "with torch.no_grad():\n",
    "    # pre-compute the image embedding\n",
    "    ts_img_embedding = sam_model.image_encoder(input_image)\n",
    "    box_torch =torch.from_numpy(np.array([[0,0,W,H]])).float().to(device)\n",
    "    if len(box_torch.shape) == 2:\n",
    "        box_torch = box_torch[:, None, :] # (B, 4) -> (B, 1, 4)\n",
    "    \n",
    "    sparse_embeddings, dense_embeddings = sam_model.prompt_encoder(\n",
    "        points=None,\n",
    "        boxes=box_torch,\n",
    "        masks=None,\n",
    "    )\n",
    "    medsam_seg_prob, _ = sam_model.mask_decoder(\n",
    "        image_embeddings=ts_img_embedding.to(device), # (B, 256, 64, 64)\n",
    "        image_pe=sam_model.prompt_encoder.get_dense_pe(), # (1, 256, 64, 64)\n",
    "        sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256)\n",
    "        dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64)\n",
    "        multimask_output=False,\n",
    "        )\n",
    "    medsam_seg_prob = torch.sigmoid(medsam_seg_prob)\n",
    "    # convert soft mask to hard mask\n",
    "    medsam_seg_prob = medsam_seg_prob.cpu().numpy().squeeze()\n",
    "    medsam_seg = (medsam_seg_prob > 0.5).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%% visualize the segmentation results of the middle slice\n",
    "# visualization functions\n",
    "# source: https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb\n",
    "# change color to avoid red and green\n",
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([251/255, 252/255, 30/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "    \n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='blue', facecolor=(0,0,0,0), lw=2))    \n",
    "\n",
    "_, axs = plt.subplots(1, 3, figsize=(25, 25))\n",
    "axs[0].imshow(image_data)\n",
    "show_mask(gt_data>0, axs[0])\n",
    "axs[0].set_title('Ground Truth', fontsize=20)\n",
    "axs[0].axis('off')\n",
    "\n",
    "axs[1].imshow(image_data)\n",
    "show_mask(ori_sam_seg, axs[1])\n",
    "axs[1].set_title('SAM Segmentation', fontsize=20)\n",
    "axs[1].axis('off')\n",
    "\n",
    "axs[2].imshow(image_data)\n",
    "show_mask(medsam_seg, axs[2])\n",
    "axs[2].set_title('MedSAM Segmentation', fontsize=20)\n",
    "axs[2].axis('off')\n",
    "plt.show()  \n",
    "plt.subplots_adjust(wspace=0.01, hspace=0)\n",
    "# save plot\n",
    "# plt.savefig(join(model_save_path, test_npzs[npz_idx].split('.npz')[0] + str(img_id).zfill(3) + '.png'), bbox_inches='tight', dpi=300)\n",
    "plt.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "medsam-demo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
